#include <stdio.h>
#include <string.h>
#include <stdarg.h>

#include "utils.h"
#include "unios.h"
#include "crc.h"
#include "ymodem.h"
#include "mcu.h"

#include "log.h"
#define LOGLEVEL LOG_INFO
#define LOGTAG   "YMODEM"
#include "log4app.h"


/// 控制字符
static const uint8_t SOH  	= 0x01; ///< Start Of Headline
static const uint8_t STX  	= 0x02; ///< Start Of Text
static const uint8_t CONT 	= 0x43; ///< 'C', Continue
static const uint8_t ACK  	= 0x06; ///< Acknowledge
static const uint8_t NAK  	= 0x15; ///< Negative Acknowledge
static const uint8_t EOT  	= 0x04; ///< End Of Transmission
static const uint8_t CAN  	= 0x18; ///< Cancel Transmission (Ctrl-X)
static const uint8_t CTRLZ	= 0x1A; ///< Blank file data (Ctrl-Z)

#define SOH_PL 128				///< 短报文 Payload 长度
#define SOH_FL (3 + SOH_PL + 2)	///< 短报文 Frame 长度
#define STX_PL 1024				///< 长报文 Payload 长度
#define STX_FL (3 + STX_PL + 2)	///< 长报文 Frame 长度

/// 中止信号
static const uint8_t CANCEL_SIGNAL[5] = { 0x18, 0x18, 0x18, 0x18, 0x18 }; // 0x18 = CAN

static const char* _statenames[] = {
	"YMODEM_STATE_IDLE",
	"YMODEM_STATE_TX_WAIT_FOR_BEGIN",		
	"YMODEM_STATE_TX_WAIT_FOR_HEAD_ACK",	
	"YMODEM_STATE_TX_WAIT_FOR_HEAD_CONT",	
	"YMODEM_STATE_TX_WAIT_FOR_DATA_ACK",	
	"YMODEM_STATE_TX_WAIT_FOR_EOT1_NAK",	
	"YMODEM_STATE_TX_WAIT_FOR_EOT2_ACK",	
	"YMODEM_STATE_TX_WAIT_FOR_EOT2_CONT",	
	"YMODEM_STATE_TX_WAIT_FOR_FIN_ACK",	
	"YMODEM_STATE_RX_WAIT_FOR_HEAD",		
	"YMODEM_STATE_RX_WAIT_FOR_DATA",		
	"YMODEM_STATE_RX_WAIT_FOR_EOT1",		
	"YMODEM_STATE_RX_WAIT_FOR_EOT2",		
	"YMODEM_STATE_RX_WAIT_FOR_FIN",		
};

static const char* _statename(ymodem_state_t state) {
	if(state >= YMODEM_STATE_MAX) return "(UNKNOWN)";
	return _statenames[state];
}

#pragma pack(push, 1)

/// 短报文格式
typedef struct {
	uint8_t type;
	uint8_t fsn;
	uint8_t fsnr;
	uint8_t data[SOH_PL];
	uint8_t crch;
	uint8_t crcl;
} soh_frame_t;

/// 长报文格式
typedef struct {
	uint8_t type;
	uint8_t fsn;
	uint8_t fsnr;
	uint8_t data[STX_PL];
	uint8_t crch;
	uint8_t crcl;
} stx_frame_t;

#pragma pack(pop)

// 简便写法: 切换协议状态机
#define _sstate(st) \
	do {							\
		logd(":"#st"\r\n");			\
		modem->state = st;			\
		_startTimeoutTimer(modem);	\
	} while(0)
	
// 简便写法: 发送数据
#define _send(data, size)	modem->sendFunc(modem->sendCookie, data, size)

static void _timeoutHandler(size_t cookie_size, void* cookie) {
	ymodem_t* modem = (ymodem_t*)mcu_get_pointer(cookie);
	logw("Timeout in state: %s !\r\n", _statename(modem->state));
	// 协议超时，中止发送
	ymodem_reset(modem, STD_ERR_TIMEOUT);
}

static void _startTimeoutTimer(ymodem_t* modem) {
	if(modem->timeoutTaskId) {
		unios_task_reset_timer(modem->timeoutTaskId);
	} else {
		modem->timeoutTaskId = unios_task_run_after(modem->TIMEOUT_MAX, _timeoutHandler, sizeof(void*), &modem);
	}
}

static void _sendFrame(ymodem_t* modem, uint8_t type) {
	// 重置重发次数
	modem->retryCount = 0;
	
	if(type == STX) {
		// 发送长报文
		stx_frame_t* frame = (stx_frame_t*)(modem->buf);
		frame->type = STX;
		frame->fsn = modem->fsn; frame->fsnr = 0xFF - modem->fsn; modem->fsn = (modem->fsn + 1) & 0xFF;
		uint16_t crc16 = crc16_xmodem(frame->data, STX_PL);
		frame->crch = HIBYTE(crc16); frame->crcl = LOBYTE(crc16);
		_send(modem->buf, sizeof(stx_frame_t));
	}
	else if(type == SOH) {
		// 发送短报文
		soh_frame_t* frame = (soh_frame_t*)(modem->buf);
		frame->type = SOH;
		frame->fsn = modem->fsn; frame->fsnr = 0xFF - modem->fsn; modem->fsn = (modem->fsn + 1) & 0xFF;
		uint16_t crc16 = crc16_xmodem(frame->data, SOH_PL);
		frame->crch = HIBYTE(crc16); frame->crcl = LOBYTE(crc16);
		_send(modem->buf, sizeof(soh_frame_t));
	}
	else if(type == 0) {
		// 发送结束帧(短报文)
		soh_frame_t* frame = (soh_frame_t*)(modem->buf);
		frame->type = SOH;
		frame->fsn = 0; frame->fsnr = 0xFF;
		memset(frame->data, 0, sizeof(frame->data)); // 数据初始化为 0
		frame->crch = 0; frame->crcl = 0;
		_send(modem->buf, sizeof(soh_frame_t));
	}
}

static std_err_t _resendFrame(ymodem_t* modem) {
	if(modem->retryCount >= modem->RETRY_COUNT_MAX) { return ymodem_reset(modem, STD_ERR_RETRY_OUT); }
	
	if(modem->buf[0] == STX) {
		_send(modem->buf, sizeof(stx_frame_t));
	} else {
		_send(modem->buf, sizeof(soh_frame_t));
	}
	
	modem->retryCount++;
	return STD_ERR_OK;
}

// 读取文件并发送数据
static std_err_t _sendFileData(ymodem_t* modem) {
	stx_frame_t* frame = (stx_frame_t*)(modem->buf);

	// 缓冲区中数据初始化为 CTRLZ
	memset(frame->data, CTRLZ, sizeof(frame->data));
	
	// 读取文件
	uint32_t size = STX_PL;
	std_err_t err = modem->reader(modem, frame->data, &size);
	if(err && err != STD_ERR_EOF) { return ymodem_reset(modem, err); }
	
	if(err == STD_ERR_EOF) {
		// 文件结束，发送 EOT
		_send(&EOT, 1);
		// 修改状态为：EOT1 已发送，等待对方 ACK/NAK
		_sstate(YMODEM_STATE_TX_WAIT_FOR_EOT1_NAK);
		return STD_ERR_OK;
	}
	
	if(size <= SOH_PL) {
		// 发送短报文
		_sendFrame(modem, SOH);
	} else {
		// 发送长报文
		_sendFrame(modem, STX);
	}

	// 修改文件已发送数据长度
	modem->fpos += size;
	// 修改状态为：文件数据已发送，等待对方 ACK
	_sstate(YMODEM_STATE_TX_WAIT_FOR_DATA_ACK);
	return STD_ERR_OK;
}

static std_err_t _receiveFrame(ymodem_t* modem, uint8_t ch) {
    if(modem->bpos == 0 || (modem->buf[0] == SOH && modem->bpos < SOH_FL) || (modem->buf[0] == STX && modem->bpos < STX_FL)) {
		// 本帧数据未接收完，数据放入缓冲区
        modem->buf[modem->bpos++] = ch;
    }
	
	if((modem->buf[0] == SOH && modem->bpos < SOH_FL) || (modem->buf[0] == STX && modem->bpos < STX_FL)) {
		// 本帧数据未接收完，继续接收
		_startTimeoutTimer(modem);
		return STD_ERR_CONTINUE;
	}
	
	// 本帧数据接收完毕
	if(modem->state == YMODEM_STATE_RX_WAIT_FOR_FIN) {
		// 结束帧，不需要判断帧序号和 CRC16
		// (sum8(modem->buf, SOH_FL) == 0x00) ? STD_ERR_OK : STD_ERR_INVALID_CRC;
	} else {
		// 检查帧序号
		if((modem->buf[1] + modem->buf[2] != 0xFF) || (modem->buf[1] != modem->fsn)) {
			// 帧序号错误，返回 NAK(重新发送本帧数据)
			loge("Invalid frame no: %d != %d\r\n", modem->buf[1], modem->fsn);
			_send(&NAK, 1);
			// 准备好重新接收本帧数据
			modem->bpos = 0;
			_startTimeoutTimer(modem);
			return STD_ERR_CONTINUE;
		}
		
		// 数据校验
		uint16_t crc16rx = MAKE_WORD(modem->buf[modem->bpos-2], modem->buf[modem->bpos-1]);
		uint16_t crc16 = crc16_xmodem(modem->buf + 3, modem->bpos - 5);
		if(crc16rx != crc16) {
			// CRC 校验错误，返回 NAK(重新发送本帧数据)
			loge("Invalid CRC 0x%04X != 0x%04X !\r\n", crc16rx, crc16);
			_send(&NAK, 1);
			// 准备好重新接收本帧数据
			modem->bpos = 0;
			_startTimeoutTimer(modem);
			return STD_ERR_CONTINUE;
		}
		
		// 数据正确，帧序号加1
		modem->fsn = (modem->fsn + 1) & 0xFF;
	}
	
	// 发回 ACK
	_send(&ACK, 1);
	return STD_ERR_OK;
}

std_err_t ymodem_on_rx_byte(ymodem_t* modem, uint8_t ch) {
	if(!modem || !modem->sendFunc) { return STD_ERR_INVALID_ARG; }
    std_err_t err = STD_ERR_OK;
	int ret;
	
	switch(modem->state) {
		default: {
			// 初始化状态
			return ymodem_reset(modem, STD_ERR_INVALID_STATE);
		}
		case YMODEM_STATE_IDLE: {
			return STD_ERR_INVALID_STATE; 
		}
        
        case YMODEM_STATE_RX_WAIT_FOR_HEAD: {
			// 接收文件数据
            err = _receiveFrame(modem, ch);
            // 当前数据帧未接收完，继续接收
            if(err == STD_ERR_CONTINUE) { break; }
            // 接收遇到错误，中止
            if(err != STD_ERR_OK) { return ymodem_reset(modem, err); }
            // 提取文件头信息
            const char* fname = (const char*)(modem->buf + 3);
            uint32_t fsize = 0;
            ret = sscanf(fname + strlen(fname) + 1, "%u", &fsize);
            if(ret != 1) { return ymodem_reset(modem, err); }
            // 处理接收到的文件头
            err = modem->header(modem, fname, fsize);
            // 文件头处理错误，中止
            if(err != STD_ERR_OK) { return ymodem_reset(modem, err); }
			// 记录当前文件大小
			modem->fsize = fsize;
			// 缓冲区数据清零
			modem->bpos = 0;
            // 发送 CONT
			_send(&CONT, 1);
			// 修改状态为：等待对端发送文件数据
			_sstate(YMODEM_STATE_RX_WAIT_FOR_DATA);
			break;
        }
        case YMODEM_STATE_RX_WAIT_FOR_DATA: {
			// 接收文件数据
            err = _receiveFrame(modem, ch);
            // 当前数据帧未接收完，继续接收
            if(err == STD_ERR_CONTINUE) { break; }
            // 接收遇到错误，中止
            if(err != STD_ERR_OK) { return ymodem_reset(modem, err); }
			// 计算有效数据长度
			int size = (modem->buf[0] == SOH) ? SOH_PL : STX_PL;
			if(modem->fpos + size >= modem->fsize) {
				// 文件已经到了末尾了，计算最后一帧有效数据长度
				size = modem->fsize - modem->fpos;
				// 修改状态为：等待对端发送 EOT1
				_sstate(YMODEM_STATE_RX_WAIT_FOR_EOT1);
			}
            // 处理接收到的文件数据
            err = modem->writer(modem, modem->buf + 3, size);
            // 文件数据处理错误，中止
            if(err != STD_ERR_OK) { return ymodem_reset(modem, err); }
			// 记录已处理的文件数据大小
			modem->fpos += size;
			// 缓冲区数据清零
			modem->bpos = 0;
			break;
        }
        case YMODEM_STATE_RX_WAIT_FOR_EOT1: {
			if(ch != EOT) { return ymodem_reset(modem, STD_ERR_INVALID_STATE); }
            // 发送 NAK
			_send(&NAK, 1);
			// 修改状态为：等待对端发送 EOT2
			_sstate(YMODEM_STATE_RX_WAIT_FOR_EOT2);
			break;
        }
        case YMODEM_STATE_RX_WAIT_FOR_EOT2: {
			if(ch != EOT) { return ymodem_reset(modem, STD_ERR_INVALID_STATE); }
            // 发送 ACK
			_send(&ACK, 1);
            // 发送 CONT
			_send(&CONT, 1);
			// 修改状态为：等待对端发送 FIN
			_sstate(YMODEM_STATE_RX_WAIT_FOR_FIN);
			break;
        }
        case YMODEM_STATE_RX_WAIT_FOR_FIN: {
			// 接收 FIN 帧
            err = _receiveFrame(modem, ch);
            // 当前数据帧未接收完，继续接收
            if(err == STD_ERR_CONTINUE) { break; }
			// 接收文件完成
			ymodem_reset(modem, 0);
			break;
        }
        
		case YMODEM_STATE_TX_WAIT_FOR_BEGIN: {
			if(ch != CONT) { return ymodem_reset(modem, STD_ERR_INVALID_STATE); }
			// 发送文件头(文件头已经准备好了，直接发送)
			// HyperTerm 应该发长报文 STX ?
			// SecureCRT 应该发短报文 SOH ?
			_sendFrame(modem, SOH);
			// 修改状态为：文件头已发送，等待对方 ACK
			_sstate(YMODEM_STATE_TX_WAIT_FOR_HEAD_ACK);
			break;
		}
		case YMODEM_STATE_TX_WAIT_FOR_HEAD_ACK: {
			// HyperTerm 会在此时回 CONT
			// SecureCRT 会在此时回 ACK
			if(ch != CONT && ch != ACK) { return ymodem_reset(modem, STD_ERR_INVALID_STATE); }
			// 修改状态为：文件头已发送，等待对方 CONT
			_sstate(YMODEM_STATE_TX_WAIT_FOR_HEAD_CONT);
			break;
		}
		case YMODEM_STATE_TX_WAIT_FOR_HEAD_CONT: {
			// HyperTerm 会在此时回 ACK
			// SecureCRT 会在此时回 CONT
			if(ch != ACK && ch != CONT) { return ymodem_reset(modem, STD_ERR_INVALID_STATE); }
			// 读取文件并发送
			_sendFileData(modem);
			break;
		}
		case YMODEM_STATE_TX_WAIT_FOR_DATA_ACK: {
			if(ch != ACK) { return ymodem_reset(modem, STD_ERR_INVALID_STATE); }
			// 读取文件并发送
			_sendFileData(modem);
			break;
		}
		case YMODEM_STATE_TX_WAIT_FOR_EOT1_NAK: 
		case YMODEM_STATE_TX_WAIT_FOR_EOT2_ACK: {
			if(modem->state == YMODEM_STATE_TX_WAIT_FOR_EOT1_NAK && ch == NAK) {
				// 需要再次发送 EOT
				_send(&EOT, 1);
				// 修改状态为：EOT2 已发送，等待对方 ACK/NAK
				_sstate(YMODEM_STATE_TX_WAIT_FOR_EOT2_ACK);
			} else if(ch == ACK) {
				// 修改状态为：EOT2 已发送，等待对方 CONT
				_sstate(YMODEM_STATE_TX_WAIT_FOR_EOT2_CONT);
			} else {
				// 协议错误，中止发送
				return ymodem_reset(modem, STD_ERR_INVALID_STATE);
			}
			break;
		}
		case YMODEM_STATE_TX_WAIT_FOR_EOT2_CONT: {
			if(ch != CONT) { return ymodem_reset(modem, STD_ERR_INVALID_STATE); }
			// 发送文件传输结束帧
			_sendFrame(modem, 0);
			// 修改状态为：文件传输结束帧已发送，等待对方 ACK
			_sstate(YMODEM_STATE_TX_WAIT_FOR_FIN_ACK);
			break;
		}
		case YMODEM_STATE_TX_WAIT_FOR_FIN_ACK: {
			if(ch != ACK) { return ymodem_reset(modem, STD_ERR_INVALID_STATE); }
			// 传输完成
			ymodem_reset(modem, 0);
			break;
		}
	}
	
	return STD_ERR_OK;
}

// 持续定时发送 CONT 的任务，以便对端启动发送文件头
static void contRepeater(size_t cookie_size, void* cookie) {
	ymodem_t* modem = (ymodem_t*)mcu_get_pointer(cookie);
	if(modem->state == YMODEM_STATE_RX_WAIT_FOR_HEAD) {
		// 发送 CONT
		_send(&CONT, 1);
	} else {
		// 停止任务
		unios_task_destroy(&modem->contRepeaterTaskId);
	}
}

std_err_t ymodem_start_receive(ymodem_t* modem) {
	if(!modem || !modem->sendFunc || !modem->header || !modem->writer) { return STD_ERR_INVALID_ARG; }
	if(modem->state != YMODEM_STATE_IDLE) { return STD_ERR_INVALID_STATE; }
	
	// 初始化状态
	ymodem_reset(modem, STD_ERR_OK);
	// 等待对方发送文件头
	_sstate(YMODEM_STATE_RX_WAIT_FOR_HEAD);
	// 启动持续定时发送 CONT 的任务，以便对端启动发送文件头
	modem->contRepeaterTaskId = unios_task_run_every(1000, modem->TIMEOUT_MAX / 1000, contRepeater, sizeof(void*), &modem, true);
	return STD_ERR_OK;
}

std_err_t ymodem_start_transmit(ymodem_t* modem, const char* fname, uint32_t fsize) {
	if(!modem || !modem->sendFunc || !modem->reader || !fname || !strlen(fname) || strlen(fname) > YMODEM_FILENAME_SIZE_MAX || !fsize) { return STD_ERR_INVALID_ARG; }
	if(modem->state != YMODEM_STATE_IDLE) { return STD_ERR_INVALID_STATE; }
	
	// 初始化状态
	ymodem_reset(modem, STD_ERR_OK);
	modem->fsize = fsize;
	
	// 准备好文件头信息
	soh_frame_t* frame = (soh_frame_t*)(modem->buf);
	
	// 数据初始化为 0
	memset(frame->data, 0, SOH_PL);
	// 填写文件名和文件长度
	sprintf((char*)frame->data, "%s %d", fname, fsize);
	// 把文件名和长度之间的空格改为 zero
	frame->data[strlen(fname)] = '\0';
	
	// 等待对方发送 CONT，以启动发送
	_sstate(YMODEM_STATE_TX_WAIT_FOR_BEGIN);
	return STD_ERR_OK;
}

std_err_t ymodem_reset(ymodem_t* modem, std_err_t errcode) {
	if(!modem) { return errcode; }
	
	if(modem->state != YMODEM_STATE_IDLE) {
		if(errcode) {
			// 发送中止信号
			if(modem->sendFunc) { _send(CANCEL_SIGNAL, sizeof(CANCEL_SIGNAL)); }
			logw("Session reset in state %s because: %s !\r\n", _statename(modem->state), std_err_comment(errcode));
		} else {
			logd("Session complete successfully !\r\n");
		}
		// 调用发送完成回调函数
		if(modem->completer) { modem->completer(modem, errcode); }
	}
	
	// 状态重置
	unios_task_destroy(&modem->timeoutTaskId);
	unios_task_destroy(&modem->contRepeaterTaskId);
	modem->state = YMODEM_STATE_IDLE;
    modem->bpos = 0;
	modem->fsn = 0;
	modem->retryCount = 0;
	modem->fsize = 0;
	modem->fpos = 0;
	modem->err = errcode;
	return errcode;
}

bool ymodem_is_transfer_completed(ymodem_t* modem) {
	return modem->state == YMODEM_STATE_IDLE;
}
